import pickle
import time
from flask import Flask, request
from concurrent.futures import ProcessPoolExecutor, as_completed
from itertools import product

app = Flask(__name__)

# Hint: 生成 psi_1t 在 T = 1, 2, 3 阶段所需要计算的全部 f_t
def psi_1t_init(v_0, v_1, v_2, v_3, epoch):
    K_t = app.config["K_t"]
    D_t = app.config["D_t"]
    psi_f_dict = {}
    for K_t_i in K_t:
        e = range(0, v_1 - v_0 + 1)
        z = range(v_3, v_3 + K_t_i + 1)
        for D_t_i in D_t:
            for _e, _z in product(e, z):
                key = f"psi_1t_f_{v_1 - _e - D_t_i}_{v_2 - _e - D_t_i}_{v_3 - _e - D_t_i}_{_z - _e - D_t_i}"
                psi_f_dict[key] = {
                    "value": 0,
                    "params": (v_1 - _e - D_t_i, v_2 - _e - D_t_i, v_3 - _e - D_t_i, _z - _e - D_t_i),
                }
    return psi_f_dict


def psi_2t_init(v_0, v_1, v_2, v_3, epoch):
    K_t = app.config["K_t"]
    D_t = app.config["D_t"]
    psi_f_dict = {}
    for K_t_i in K_t:
        y_12 = range(v_1, v_2 + 1)
        z = range(v_3, v_3 + K_t_i + 1)
        for D_t_i in D_t:
            for _y_12, _z in product(y_12, z):
                key = f"psi_2t_f_{_y_12 - D_t_i}_{v_2 - D_t_i}_{v_3 - D_t_i}_{_z - D_t_i}"
                psi_f_dict[key] = {
                    "value": 0,
                    "params": (_y_12 - D_t_i, v_2 - D_t_i, v_3 - D_t_i, _z - D_t_i),
                }
    return psi_f_dict


def psi_3t_init(v_0, v_1, v_2, v_3, epoch):
    K_t = app.config["K_t"]
    D_t = app.config["D_t"]
    psi_f_dict = {}
    for K_t_i in K_t:
        y_23 = range(v_2, v_3 + 1)
        z = range(v_3, v_3 + K_t_i + 1)
        for D_t_i in D_t:
            for _y_23, _z in product(y_23, z):
                key = f"psi_3t_f_{_y_23 - D_t_i}_{_y_23 - D_t_i}_{v_3 - D_t_i}_{_z - D_t_i}"
                psi_f_dict[key] = {
                    "value": 0,
                    "params": (_y_23 - D_t_i, _y_23 - D_t_i, v_3 - D_t_i, _z - D_t_i),
                }
    return psi_f_dict


def grid_search_exa(iter_grid, func, params: tuple) -> list:
    pool = app.config["pool"]
    tasks = [pool.submit(func, *params, x, y) for x, y in iter_grid]
    results = [t.result() for t in as_completed(tasks)]
    return results


def linear_search_exa(iter_arr, func, params: tuple) -> list:
    pool = app.config["pool"]
    tasks = [pool.submit(func, *params, x) for x in iter_arr]
    results = [t.result() for t in as_completed(tasks)]
    return results


@app.route("/")
def r():
    return "This is backstage server for compute service."


@app.route("/init_service")
def r_init_service():
    if app.config.get("pool") is not None:
        return {"message": "pool is already initialized"}
    try:
        data = request.get_json()
        pool = ProcessPoolExecutor(max_workers=data["max_worker"])
    except Exception as e:
        print("not given max_worker param, use default max_worker")
        pool = ProcessPoolExecutor()
    app.config["pool"] = pool
    return {"message": "init service success"}


@app.route("/init_model", methods=["POST"])
def r_init_model():
    if app.config.get("global_cache") is not None:
        return {"message": "global_cache is already initialized"}
    try:
        data = pickle.loads(request.get_data())
        T = data["T"]
        v_0 = data["v_0"]
        v_1 = data["v_1"]
        v_2 = data["v_2"]
        v_3 = data["v_3"]
        alpha = data["alpha"]
        p = data["p"]
        b = data["b"]
        c_e = data["c_e"]
        c_m = data["c_m"]
        h_e = data["h_e"]
        h_m = data["h_m"]
        theta = data["theta"]
        K_t = data["K_t"]
        p_K_t = data["p_K_t"]
        D_t = data["D_t"]
        p_D_t= data["p_D_t"]

        app.config["T"] = T
        app.config["alpha"] = alpha
        app.config["p"] = p
        app.config["b"] = b
        app.config["c_e"] = c_e
        app.config["c_m"] = c_m
        app.config["h_e"] = h_e
        app.config["h_m"] = h_m
        app.config["theta"] = theta
        app.config["K_t"] = K_t
        app.config["p_K_t"] = p_K_t
        app.config["D_t"] = D_t
        app.config["p_D_t"] = p_D_t

        ### 初始化 global_cache

        init_begin_time = time.perf_counter()

        epoch = 1
        max_epoch = T - 1
        global_cache = {}
        global_cache[epoch] = {}
        global_cache[epoch].update(psi_1t_init(v_0, v_1, v_2, v_3, epoch))
        global_cache[epoch].update(psi_2t_init(v_0, v_1, v_2, v_3, epoch))
        global_cache[epoch].update(psi_3t_init(v_0, v_1, v_2, v_3, epoch))

        while True:
            if epoch == max_epoch:
                break

            global_cache[epoch + 1] = {}
            # 遍历上一轮所有需要计算的 f_t
            for key in global_cache[epoch].keys():
                params = global_cache[epoch][key]["params"]
                global_cache[epoch + 1].update(psi_1t_init(*params, epoch))
                global_cache[epoch + 1].update(psi_2t_init(*params, epoch))
                global_cache[epoch + 1].update(psi_3t_init(*params, epoch))
            epoch = epoch + 1

        init_end_time = time.perf_counter()
        print(f"模型初始化完成. 耗时 {init_end_time - init_begin_time: .2f} 秒")

        ### 初始化 global_cache 完成

        app.config["global_cache"] = global_cache

        ### 逐轮计算 global_cache



    except Exception as e:
        print(e.with_traceback())
        return {"message": "init failed"}

    return {"message": "init success"}


@app.route("/grid_search_exa", methods=["POST"])
def r_grid_search_exa():
    return {}


if __name__ == "__main__":
    app.run(debug=True, port=5000)
